from pyzerox import zerox as original_zerox
import os
import aiofiles
import asyncio
import sys
import io

async def custom_zerox(file_path, model, output_dir="./output", custom_system_prompt=None, select_pages=None, **kwargs):
    """
    自定义的 zerox 函数，使用 UTF-8 编码保存结果
    """
    # 确保标准输出使用UTF-8编码
    if sys.stdout.encoding != 'utf-8':
        try:
            sys.stdout.reconfigure(encoding='utf-8')
            sys.stderr.reconfigure(encoding='utf-8')
        except AttributeError:
            sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
            sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')
    
    # 调用原始 zerox 函数获取结果，但不让它保存文件
    temp_output_dir = "./temp_output_" + str(int(asyncio.get_event_loop().time() * 1000))
    os.makedirs(temp_output_dir, exist_ok=True)
    
    try:
        result = await original_zerox(
            file_path=file_path,
            model=model,
            output_dir=temp_output_dir,
            custom_system_prompt=custom_system_prompt,
            select_pages=select_pages,
            **kwargs
        )
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 获取文件名（不含扩展名）
        file_name = os.path.basename(file_path)
        file_name = os.path.splitext(file_name)[0]
        
        # 构建输出文件路径
        output_file = os.path.join(output_dir, f"{file_name}.md")
        
        # 使用同步方式写入文件，避免 aiofiles 可能的编码问题
        if isinstance(result, str):
            try:
                # 尝试以 UTF-8 编码写入文件，忽略编码错误
                with open(output_file, 'w', encoding='utf-8', errors='ignore') as f:
                    f.write(result)
                print(f"结果已保存到: {output_file}")
            except Exception as e:
                print(f"写入文件时出错: {str(e)}")
                # 如果仍然失败，尝试备用方法
                try:
                    import codecs
                    with codecs.open(output_file, 'w', encoding='utf-8', errors='ignore') as f:
                        f.write(result)
                    print(f"使用 codecs 成功保存结果到: {output_file}")
                except Exception as e2:
                    print(f"备用写入方法也失败: {str(e2)}")
        else:
            # 如果结果是对象，尝试提取内容
            try:
                content = str(result)
                with open(output_file, 'w', encoding='utf-8', errors='ignore') as f:
                    f.write(content)
            except Exception as e:
                print(f"处理结果格式错误: {str(e)}")
        
        return result
    finally:
        # 清理临时目录
        import shutil
        try:
            shutil.rmtree(temp_output_dir)
        except:
            pass 